package com.ttt.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

@RestController
@RequestMapping("/ai")
public class AIBaseController {
    @Autowired
    private ChatClient chatClient;
    @Autowired
    private DocumentRetriever documentRetriever;
    /**
     * 知识库问答接口
     */
    @PostMapping("/tast1")
    public String tast1(String question) {
        return chatClient.prompt()
                .user(question)
                .advisors(RetrievalAugmentationAdvisor.builder().documentRetriever(documentRetriever).queryExpander(new MultiQueryExpander(chatClient.mutate(),
                        new PromptTemplate("您是信息检索和搜索优化方面的专家。（主要检索PCB方案方面知识）\n" +
                                "您的任务是生成给定查询的 {number} 个不同版本。\n" +
                                "\n" +
                                "每个变体必须涵盖主题的不同观点或多个方面，（比如：'客户查询PCB的知识'你可以扩展到PCB检测、PCB设计、PCB项目计划、PCB技术等多个方面）\n" +
                                "同时保持原始查询的核心意图。目标是\n" +
                                "扩大搜索空间并提高找到相关信息的机会。（你可以尽可能地扩大搜索空间，但不要过度，也不要包含太多无关原查询领域的内容）\n" +
                                "\n" +
                                "请勿解释您的选择或添加任何其他无助与查询的文本。\n" +
                                "提供以换行符分隔的查询变体。你可以在用户的原始查询上进行适当的扩展查询内容（一定是适当，而不是过度）\n" +
                                "\n" +
                                "原始查询： {query}\n" +
                                "\n" +
                                "查询变体：\n" +
                                "\"\"\");"),true,3)).build())
                .call()
                .content();
    }
}
